#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import re
import io
import socket
import shutil
import pycdlib
import gstage4
import wstage4
import pathlib
import tempfile
import subprocess
import grub_install
import strict_portage
import multiprocessing
from fm_util import Util
from fm_util import CloudCacheGentoo
from fm_util import CcacheLocalService
from fm_util import PrintLoadAvgThread
from fm_util import TmpMount
from fm_param import FmConst


class AssistOsBuilder:

    def __init__(self, infoPrinter, cloudOverlayDb, machine, bHasMemTest, showName):
        self._targetSystemInfo = {
            "amd64": {
                "subarch": "amd64",
                "variant": "systemd",
                "profile": "default/linux/amd64/23.0/no-multilib",
                "stage3-file": None,
            },
            "arm64": {
                "subarch": "arm64",
                "variant": "systemd",
                "profile": None,
                "stage3-file": None,
            },
        }
        self._p = infoPrinter
        self._cloudOverlayDb = cloudOverlayDb
        self._machine = machine
        self._bHasMemTest = bHasMemTest
        self._showName = showName
        self._snapshotFile = None

    def downloadFiles(self, bTitle):
        if bTitle:
            self._p.printInfo("- Downloading %s files..." % (self._showName))

        cache = CloudCacheGentoo(FmConst.gentooLinuxCacheDir)
        cache.sync()
        for arch, v in self._targetSystemInfo.items():
            assert arch in cache.get_arch_list()
            assert v["subarch"] in cache.get_subarch_list(arch)
        for arch, v in self._targetSystemInfo.items():
            try:
                v["stage3-file"] = cache.get_latest_stage3(arch, v["subarch"], v["variant"], cached_only=True)   # prefer local stage3 file
            except FileNotFoundError:
                v["stage3-file"] = cache.get_latest_stage3(arch, v["subarch"], v["variant"])
        self._snapshotFile = cache.get_latest_snapshot()                                                         # always use newest snapshot

    def build(self, arch, diskFsUuid, wdir):
        ftFixbug = gstage4.target_features.FixBugs()
        ftKernel = gstage4.target_features.UseBbki()
        ftPackageManager = gstage4.target_features.UsePortage()
        ftServiceManager = gstage4.target_features.UseSystemd()
        ftNoDeprecate = gstage4.target_features.DontUseDeprecatedPackagesAndFunctions()
        ftUseCapability = gstage4.target_features.UseCapability()
        ftPreferGnu = gstage4.target_features.PreferGnuAndGpl()
        ftPreferBinPkg = gstage4.target_features.PreferBinaryPackage()
        ftPreferSystemComponent = gstage4.target_features.PreferSystemComponent()
        ftPreferWayland = gstage4.target_features.PreferWayland()
        ftPreferBlockDeviceUAccess = gstage4.target_features.PreferBlockDeviceUAccess()
        ftPreferPythonMagic = gstage4.target_features.PreferPythonMagic()
        # ftNoVt = gstage4.target_features.DontUseVT()
        ftDisableFstab = gstage4.target_features.DisableFstab()
        ftNotUsePypy = gstage4.target_features.DontUsePypy()
        ftSupportAllTermType = gstage4.target_features.SupportAllTermType()
        # ftKmscon = gstage4.target_features.Kmscon()
        ftSshServer = gstage4.target_features.SshServer()
        if self._bHasMemTest:
            ftMemTest = gstage4.target_features.MemTest()
        ftChrony = gstage4.target_features.Chrony(exclusive=True)
        ftNetworkManager = gstage4.target_features.NetworkManager(wifi=True, exclusive=True)
        ftZeroConf = gstage4.target_features.UseZeroConf(service="avahi")
        ftNoPolkit = gstage4.target_features.DontUsePolicyKit()
        ftNoLogrotate = gstage4.target_features.DontUseLogrotate()
        ftNoSudo = gstage4.target_features.DontUseSudo()
        ftTailorSystemd = gstage4.target_features.TailorSystemd(
            disable_items=[
                "systemd-udevd-socket-activation",
                "kmod-static-nodes",
            ],
            remove_items=[
                "systemd-battery-check",
                "systemd-boot",
                "systemd-coredump",                         # diagnosing things are beyound us
                "systemd-dissect",
                "systemd-hostnamed",                        # we don't change hostname dynamically
                "systemd-firstboot",
                "systemd-kexec",
                "systemd-localed",                          # eselect is enough
                "systemd-machined",
                "systemd-networkd",
                "systemd-oomd",
                "systemd-portabled",
                "systemd-pstore",
                "systemd-resolvd",
                "systemd-storagetm",
                "systemd-sysext",
                "systemd-sysupdate",
                "systemd-sysusers",
                "systemd-timedated",                        # eselect is enough
                "systemd-timesyncd",                        # chrony is better
                "systemd-update",
                "systemd-userdbd",
                "systemd-ssh-proxy",
                "systemd-run-generator",                    # we don't tell systemd to run something through kernel command line
                "systemd-efi-boot-generator",               # we mount /boot in initramfs so this is uneccessary
                "systemd-gpt-auto-generator",               # all the key mount point are established in initramfs, no automation needed
                "systemd-ssh-generator",
                "kernel-management",
                "fstab",
                "fs-operations",
                "ldconfig.service",                         # we update ld cache during package install, not system boot
                "initrd-facility",                          # we don't use systemd in initramfs
                "debug-facility",
            ]
        )
        ftTailorEselect = gstage4.target_features.TailorEselect(remove_items=[
            "kernel-module",
            "rc-module",
        ])
        ftTailorGit = gstage4.target_features.TailorGit(add_items=[
            "http-connection-timeout-extension",
            "robust-extension",
        ])
        ftTailorWget = gstage4.target_features.TailorWget(add_items=[
            "robust-extension",
        ])
        ftTailorRsync = gstage4.target_features.TailorRsync(add_items=[
            "robust-extension",
        ])
        ftAssistOs = TargetFeatureAssistOs(self._machine)

        # step
        self._p.printInfo("- Initializing...")
        c = CcacheLocalService()
        builder = None
        if True:
            s = gstage4.Settings()
            s.program_name = FmConst.programName
            s.log_dir = None
            s.verbose_level = 0
            s.host_cpu_core_count = multiprocessing.cpu_count()
            s.host_memory_size = Util.getPhysicalMemorySize() * 1024 * 1024 * 1024
            s.host_cooling_level = 10 if Util.hwcfgHasFan(self._machine) is not None else 1
            s.host_distfiles_dir = FmConst.distDir
            if c.is_enabled():
                s.host_ccache_dir = c.get_ccache_dir()

            ts = gstage4.TargetSettings()
            ts.arch = arch
            ts.profile = self._targetSystemInfo[arch]["profile"]
            ts.mirrors = strict_portage.cfg.PortageConfigDir().get_mirrors_obj().get_mirror_mapping()
            ftFixbug.update_target_settings(ts)
            ftKernel.update_target_settings(ts)
            ftPackageManager.update_target_settings(ts)
            ftServiceManager.update_target_settings(ts)
            ftNoDeprecate.update_target_settings(ts)
            ftUseCapability.update_target_settings(ts)
            ftPreferGnu.update_target_settings(ts)
            ftPreferBinPkg.update_target_settings(ts)
            ftPreferSystemComponent.update_target_settings(ts)
            ftPreferWayland.update_target_settings(ts)
            ftPreferBlockDeviceUAccess.update_target_settings(ts)
            ftPreferPythonMagic.update_target_settings(ts)
            # ftNoVt.update_target_settings(ts)
            ftDisableFstab.update_target_settings(ts)
            ftNotUsePypy.update_target_settings(ts)
            if self._bHasMemTest:
                ftMemTest.update_target_settings(ts)
            # ftKmscon.update_target_settings(ts)
            ftChrony.update_target_settings(ts)
            ftNetworkManager.update_target_settings(ts)
            ftNoPolkit.update_target_settings(ts)
            ftNoLogrotate.update_target_settings(ts)
            ftNoSudo.update_target_settings(ts)
            ftTailorSystemd.update_target_settings(ts)
            ftTailorEselect.update_target_settings(ts)
            ftTailorGit.update_target_settings(ts)
            ftTailorWget.update_target_settings(ts)
            ftTailorRsync.update_target_settings(ts)
            ftAssistOs.update_target_settings(ts)
            if c.is_enabled():
                # we use ccache if host uses ccache
                ts.build_opts.ccache = True

            builder = gstage4.Builder(s, ts, wdir)

        # step
        self._p.printInfo("- Extracting seed stage...")
        with gstage4.seed_stages.GentooStage3Archive(*self._targetSystemInfo[arch]["stage3-file"]) as ss:
            builder.action_unpack(ss)

        # step
        self._p.printInfo("- Installing gentoo repository...")
        gentooRepo = gstage4.repositories.GentooSnapshot(self._snapshotFile)
        builder.action_create_gentoo_repository(gentooRepo)

        # step
        self._p.printInfo("- Generating configurations...")
        if True:
            builder.action_init_confdir()

        # step
        self._p.printInfo("- Installing overlays...")
        builder.action_create_overlays([
            gstage4.repositories.UserDefinedOverlay("guru", *self._cloudOverlayDb.getOverlayVcsTypeAndUrl("guru")),
            gstage4.repositories.UserDefinedOverlay("mirrorshq-overlay", "git", "https://gitee.com/mirrorshq/gentoo-overlay"),
            gstage4.repositories.UserDefinedOverlay("your-own-os-overlay", "git", "https://gitee.com/your-own-os/gentoo-overlay"),
            gstage4.repositories.UserDefinedOverlay("fpemud-overlay", "git", "https://gitee.com/fpemud/gentoo-overlay"),                         # app-admin/fpemud-os-assist-disk-scripts
        ])

        # step
        with PrintLoadAvgThread(self._p.getIndentStr() + "- Updating world..."):
            worldSet = {
                "app-admin/barkspider",
                "app-admin/eselect",
                "app-admin/fpemud-os-assist-disk-scripts",
                "app-arch/cpio",
                "app-arch/gzip",
                "app-arch/lz4",
                "app-arch/p7zip",
                "app-arch/unzip",
                "app-arch/zip",
                "app-eselect/eselect-timezone",
                "app-editors/nano",
                "app-misc/mc",
                "app-misc/tmux",
                "app-text/tree",
                "dev-debug/strace",
                "dev-lang/python",
                "dev-python/bbki",
                "dev-vcs/git",
                "dev-vcs/subversion",
                "gui-wm/cage",
                "net-misc/rsync",
                "net-misc/wget",
                "sys-apps/dmidecode",
                "sys-apps/lm-sensors",
                "sys-apps/lshw",
                "sys-apps/smartmontools",
                "sys-boot/grub",
                "sys-apps/file",
                "sys-apps/hdparm",
                "sys-apps/nvme-cli",            # manipulate nvme device
                "sys-apps/sdparm",              # manipulate sata device
                "sys-block/parted",
                "sys-boot/os-prober",
                "sys-devel/bc",
                "sys-fs/bcache-tools",          # manipulate bcache
                "sys-fs/btrfs-progs",           # manipulate btrfs filesystem
                "sys-fs/dosfstools",            # manipulate fat filesystem
                "sys-fs/e2fsprogs",             # manipulate ext2/3/4 filesystem
                "sys-fs/exfatprogs",            # manipulate exfat filesystem
                "sys-fs/f2fs-tools",            # manipulate f2fs filesystem
                "sys-fs/lsscsi",
                "sys-fs/mtools",
                "sys-fs/ntfs3g",                # manipulate ntfs filesystem
                "sys-process/bpytop",
                "sys-process/lsof",
                "sys-process/iotop",
                "x11-terms/kitty",
            }
            if c.is_enabled():
                worldSet.add("dev-util/ccache")
            ftPackageManager.update_world_set(worldSet)
            ftServiceManager.update_world_set(worldSet)
            ftSupportAllTermType.update_world_set(worldSet)
            # ftKmscon.update_world_set(worldSet)
            ftSshServer.update_world_set(worldSet)
            if self._bHasMemTest:
                ftMemTest.update_world_set(worldSet)
            ftChrony.update_world_set(worldSet)
            ftNetworkManager.update_world_set(worldSet)
            ftZeroConf.update_world_set(worldSet)

            builder.action_update_world(worldSet)

            builder.add_and_run_custom_action("robustify_by_barkspider", ftAssistOs.get_custom_action_robustify_by_barkspider())

        # step
        with PrintLoadAvgThread(self._p.getIndentStr() + "- Building kernel..."):
            builder.action_install_kernel()     # do nothing
            builder.add_and_run_custom_action("real_build_kernel_and_initramfs", ftAssistOs.get_custom_action_build_kernel_and_initramfs(arch, diskFsUuid, self._p.getIndentStr()))

        # step
        self._p.printInfo("- Enabling services...")
        if True:
            serviceList = []
            # ftKmscon.update_service_list(serviceList)
            ftSshServer.update_service_list(serviceList)
            ftChrony.update_service_list(serviceList)
            ftNetworkManager.update_service_list(serviceList)
            ftZeroConf.update_service_list(serviceList)

            builder.action_enable_services(serviceList)

        # step
        self._p.printInfo("- Customizing...")
        if True:
            builder.add_and_run_custom_action("set_root_password", gstage4.custom_actions.SetRootPassword("123456"))
            builder.add_and_run_custom_action("remove_usr_src_dir_content", gstage4.custom_actions.RemoveUsrSrcDirectoryContent())
            builder.add_and_run_custom_action("generate_machine_id", ftAssistOs.get_custom_action_generate_machine_id())
            builder.add_and_run_custom_action("enable_console_service", ftAssistOs.get_custom_action_enable_console_service())

        # step
        self._p.printInfo("- Cleaning up...")
        builder.action_cleanup()

        # finish
        builder.finish()

    def updateInitramfs(self, diskFsUuid, wdir):
        buf = """
            import bbki

            class MyConfig(bbki.ConfigBase):

                @property
                def flags(self):
                    return bbki.Flags.KERNEL | bbki.Flags.INITRAMFS

                def get_fs_layout(self):
                    return bbki.FsLayout()

                def get_build_variable(self, var_name):
                    return ""

                def get_mount_points(self):
                    return [
                        bbki.HostMountPoint(**{
                            "mountpoint": "/",
                            "dev_path_or_uuid": "UUID:FILE=%s:os/%s/rootfs.sqfs",
                            "fstype": "squashfs+overlay",
                            "opts": "",
                            "underlay_disk": {
                                "type": "image-file-on-removable",
                                "uuid": "UUID=%s",
                                "used_partition_table_types": ["mbr", "gpt"],
                                "image_file_path": "os/%s/rootfs.sqfs",
                                "children": [],
                            },
                        }),
                    ]

                def get_system_init(self):
                    return bbki.SystemInit(bbki.SystemInit.Type.SYSTEMD)

                def get_online_mount_points(self):
                    assert False

                def get_boot_mode(self):
                    assert False

                def get_bootloader_extra_time(self):
                    assert False

                def get_kernel_extra_init_cmdline(self):
                    assert False

            bbkiObj = bbki.Bbki(MyConfig())
            updater = bbkiObj.obtain_kernel_updater(bbkiObj.get_newest_boot_entry())
            updater.update_initramfs()
        """
        buf = re.sub("^            ", "", buf, flags=re.M)
        buf = buf % (diskFsUuid, wdir.get_arch(), diskFsUuid, wdir.get_arch())

        fullfn = os.path.join(wdir.get_latest_action_dirpath(), "root", "update-initramfs")
        pathlib.Path(fullfn).write_text(buf)
        with gstage4.Runner(wdir.get_arch(), wdir.get_latest_action_dirpath()) as c:
            c.shell_call("", "python /root/update-initramfs")
        os.unlink(fullfn)


class TargetFeatureAssistOs:

    def __init__(self, machine):
        self._machine = machine

    def update_target_settings(self, target_settings):
        assert "80-assist-os" not in target_settings.pkg_use_files

        target_settings.pkg_use_files["80-assist-os"] = self._useFileContent.strip("\n") + "\n"

    def get_custom_action_robustify_by_barkspider(self):
        custom_script_list = []

        # re-install packages so that they can be robustified
        custom_script_list.append(gstage4.scripts.OneLinerScript("emerge -1 dev-vcs/git net-misc/rsync net-misc/wget"))

        # change configure so that the applications can be robustified
        custom_script_list.append(gstage4.scripts.OneLinerScript("barkspider robustify all"))

        return gstage4.custom_actions.SimpleCustomAction(*custom_script_list)

    def get_custom_action_build_kernel_and_initramfs(self, arch, diskFsUuid, indentStr):
        jobcountMake, loadavg = Util.portageDetectJobCountAndLoadAverage(self._machine)

        buf = """
            import bbki
            import platform

            class MyConfig(bbki.ConfigBase):

                @property
                def flags(self):
                    return bbki.Flags.KERNEL | bbki.Flags.INITRAMFS

                def get_fs_layout(self):
                    return bbki.FsLayout()

                def get_build_variable(self, var_name):
                    return "--jobs=%d --load-average=%d"

                def get_mount_points(self):
                    return [
                        bbki.HostMountPoint(**{
                            "mountpoint": "/",
                            "dev_path_or_uuid": "UUID:FILE=%s:os/%s/rootfs.sqfs",
                            "fstype": "squashfs+overlay",
                            "opts": "",
                            "underlay_disk": {
                                "type": "image-file-on-removable",
                                "uuid": "UUID=%s",
                                "used_partition_table_types": ["mbr", "gpt"],
                                "image_file_path": "os/%s/rootfs.sqfs",
                                "children": [],
                            },
                        }),
                    ]

                def get_system_init(self):
                    return bbki.SystemInit(bbki.SystemInit.Type.SYSTEMD)

                def get_online_mount_points(self):
                    assert False

                def get_boot_mode(self):
                    assert False

                def get_bootloader_extra_time(self):
                    assert False

                def get_kernel_extra_init_cmdline(self):
                    assert False

            indentStr = "%s"

            bbkiObj = bbki.Bbki(MyConfig())

            bootEntrySpec = bbki.BootEntrySpec({
                "arch": platform.machine(),
                "kernel_type": "linux",
                "kernel_params": {
                },
                "kernel_addon_params": {
                    "pnp": {
                        "peripheral": None,
                    },
                    "noatime_default": {
                    },
                    "wireless_regdb": {
                    },
                    "linux_firmware": {
                        "strict": None
                    },
                    "fedora_alsa_firmware": {
                        "strict": None
                    },
                    "riptide_firmware": {
                        "strict": None
                    },
                    "ipw2100_firmware": {
                        "strict": None
                    },
                    "ipw2200_firmware": {
                        "strict": None
                    },
                    "storage_layouts": {
                        "efi-bcache-btrfs": None,       # for target os
                        "efi-btrfs": None,              # for target os
                        "efi-ext4": None,               # for target os
                        "efi-msr-ntfs": None,           # for target os
                        "bios-ext4": None,              # for target os
                        "bios-ntfs": None,              # for target os
                        "squashfs-on-fat": None,
                        "squashfs-on-isofs": None,
                    },
                    "x86_32": {
                    },
                    "x86_intel": {
                    },
                    "x86_amd": {
                    },
                    "iwd": {
                    },
                    "vm_as_host": {
                    },
                    "vt": {
                    },
                    # "display_full": {
                    # },
                    # "display_simple": {
                    #     "boot-logger": None,
                    # },
                    "desktop_functions": {
                    },
                },
                "initramfs_params": {
                },
            })

            print(indentStr, end="")
            print("        - Synchronizing...")
            bbkiObj.sync(bootEntrySpec)

            print(indentStr, end="")
            print("        - Fetching...")
            bbkiKernelBuilder = bbkiObj.obtain_kernel_installer(bootEntrySpec)
            bbkiKernelBuilder.fetch_kernel()
            for name in bbkiKernelBuilder.spec.kernel_addons:
                bbkiKernelBuilder.fetch_kernel_addon(name)

            print(indentStr, end="")
            print("        - Extracting...")
            bbkiKernelBuilder.unpack()

            print(indentStr, end="")
            print("        - Patching...")
            bbkiKernelBuilder.patch_kernel()

            print(indentStr, end="")
            print("        - Generating .config file...")
            bbkiKernelBuilder.generate_kernel_config_file()

            print(indentStr, end="")
            print("        - Building and installing...")
            bbkiKernelBuilder.install_kernel()

            print(indentStr, end="")
            print("        - Creating initramfs...")
            bbkiKernelBuilder.install_initramfs()
        """
        buf = re.sub("^            ", "", buf, flags=re.M)
        buf = buf % (jobcountMake, loadavg, diskFsUuid, arch, diskFsUuid, arch, indentStr)

        custom_script_list = []

        s = gstage4.scripts.PlacingFilesScript()
        s.append_file("/root/build-kernel", buf)
        custom_script_list.append(s)

        custom_script_list.append(gstage4.scripts.OneLinerScript("python /root/build-kernel"))

        return gstage4.custom_actions.SimpleCustomAction(*custom_script_list)

    def get_custom_action_generate_machine_id(self):
        s = gstage4.scripts.PlacingFilesScript()
        s.append_file("/etc/machine-id", Util.generateRandomMachineId() + "\n")
        return gstage4.custom_actions.SimpleCustomAction(s)

    # def get_custom_action_establish_seat_main(self):
    #     s = gstage4.scripts.PlacingFilesScript()
    #     s.append_file("/etc/udev/rules.d/72-seat-main.rules", 'TAG=="seat", ENV{ID_SEAT}="seat-main"\n')
    #     return gstage4.custom_actions.SimpleCustomAction(s)

    def get_custom_action_enable_console_service(self):
        buf = ""
        buf += "[Unit]\n"
        buf += "Description=Assistant Disk Console\n"
        buf += "\n"
        buf += "[Service]\n"
        buf += "ExecStart=/usr/bin/cage -- /usr/bin/kitty\n"
        buf += "\n"
        buf += "[Install]\n"
        buf += "WantedBy=graphical.target\n"

        s = gstage4.scripts.PlacingFilesScript()
        s.append_file("/etc/systemd/system/assist-disk-console.service", buf)
        return gstage4.custom_actions.SimpleCustomAction(s)

    def get_custom_action_create_livecd_mount_point(self):
        assert False

    _useFileContent = """
"""


class AssistDiskBuilder:

    def __init__(self, diskFsUuid, diskLabel, diskName):
        self._diskFsUuid = diskFsUuid
        self._diskLabel = diskLabel
        self._diskName = diskName
        self._wdirStage4GentooList = []
        self._wdirStage4WinXp = None
        self._wdirStage4Win7 = None
        self._targetSysDict = {}            # dict<arch, wdir>

    def addStage4(self, osType, wdir):
        if osType == "gentoo":
            assert isinstance(wdir, gstage4.WorkDir)
            assert wdir.is_build_finished()
            assert wdir.get_arch() not in [x.get_arch() for x in self._wdirStage4GentooList]
            self._wdirStage4GentooList.append(wdir)
        elif osType == "mswin_xp":
            assert isinstance(wdir, wstage4.WorkDir)
            assert wdir.is_build_finished()
            assert self._wdirStage4WinXp is None
            self._wdirStage4WinXp = wdir
        elif osType == "mswin_7":
            assert isinstance(wdir, wstage4.WorkDir)
            assert wdir.is_build_finished()
            assert self._wdirStage4Win7 is None
            self._wdirStage4Win7 = wdir
        else:
            assert False

    def addAssistOs(self, wdir):
        assert isinstance(wdir, gstage4.WorkDir) and wdir.is_build_finished()
        assert wdir.get_arch() not in self._targetSysDict
        self._targetSysDict[wdir.get_arch()] = wdir

    def exportToIsoFile(self, filePath):
        assert len(self._targetSysDict) > 0

        iso = pycdlib.PyCdlib()
        iso.new(udf="2.60")
        try:
            # add README.TXT
            buf = ""
            buf += 'This disc contains a "UDF" file system and requires an operating system\n'
            buf += 'that supports the ISO-13346 "UDF" file system specification.\n'
            buf = buf.encode("iso8859-1")
            iso.add_fp(io.ByteIO(buf), len(buf), iso_path="/README.TXT")

            # add files
            f = iso.get_udf_facade()
            f.add_directory("/os")
            f.add_directory("/data")

            # add boot files
            pass

            # write
            iso.write(filePath)
        finally:
            iso.close()

    def exportToDiskImageFile(self, filePath, imageSize):
        tDir = tempfile.mkdtemp()
        try:
            mntDir = os.path.join(tDir, "mnt")
            os.mkdir(mntDir)
            loopFile = os.path.join(tDir, "loop")

            Util.createAndMountFatDiskImage(filePath, imageSize, "mbr", loopFile, mntDir, fsUuid=self._diskFsUuid, label=self._diskLabel)
            try:
                mp = grub_install.RootfsOrBootMountPoint(filePath, loopFile, mntDir, "vfat", "", False, True)
                self._exportToDir(mp)
            finally:
                subprocess.check_call(["umount", mntDir])
                Util.loopUnMount(loopFile)
        finally:
            shutil.rmtree(tDir)

    def exportToUsbStick(self, devPath):
        assert "amd64" in self._targetSysDict or "arm64" in self._targetSysDict

        partDevPath = Util.formatDisk(devPath, "mbr", fsType="fat", fsUuid=self._diskFsUuid, label=self._diskLabel)
        with TmpMount(partDevPath) as mp:
            mp = grub_install.RootfsOrBootMountPoint(devPath, partDevPath, mp.mountpoint, mp.fstype, mp.opts, False, True)
            self._exportToDir(mp)

    def _exportToDir(self, mp):
        uuid = Util.getBlkDevUuid(mp.partition)
        if uuid == "":
            raise Exception("can not get FS-UUID for %s" % (mp.partition))

        osDir = os.path.join(mp.mountpoint, "os")
        os.mkdir(osDir)
        dataDir = os.path.join(mp.mountpoint, "data")
        os.mkdir(dataDir)

        # copy gentoo stage4 files
        for wdir in self._wdirStage4GentooList:
            # each volume must have size of 4G because we are using vfat file system
            dstFile = os.path.join(dataDir, "gentoo-%s.tar.xz" % (wdir.get_arch()))
            Util.compressMultiVolumeTarXz(wdir.get_latest_action_dirpath(), dstFile, "2G")    # FIXME: 4G repports too large error, why?

        # copy mswin_xp stage4 files
        if self._wdirStage4WinXp is not None:
            with self._wdirStage4WinXp.mount_image_file() as m:
                with open(os.path.join(dataDir, "mswin_xp-mbr_boot_code.bin"), "wb") as f:    # read boot code from master boot sector
                    f.write(Util.mbrReadBootCode(self._wdirStage4WinXp.image_filepath))
                with open(os.path.join(dataDir, "mswin_xp-vbr_boot_code.bin"), "wb") as f:    # read boot code from volume boot sector
                    f.write(Util.ntfsReadBootCode(m.drive_c_loop_file))
                Util.compressMultiVolume7z(m.drive_c_dir, os.path.join(dataDir, "mswin_xp-drive_c.7z"), "4G")

        # copy mswin_7 stage4 files
        if self._wdirStage4Win7 is not None:
            with self._wdirStage4Win7.mount_image_file() as m:
                with open(os.path.join(dataDir, "mswin_7-mbr_boot_code.bin"), "wb") as f:    # read boot code from master boot sector
                    f.write(Util.mbrReadBootCode(self._wdirStage4Win7.image_filepath))
                with open(os.path.join(dataDir, "mswin_7-vbr_boot_code.bin"), "wb") as f:    # read boot code from volume boot sector
                    f.write(Util.ntfsReadBootCode(m.drive_c_loop_file))
                Util.compressMultiVolume7z(m.drive_c_dir, os.path.join(dataDir, "mswin_7-drive_c.7z"), "4G")

        # install target system files (amd64) into usb stick
        if "amd64" in self._targetSysDict:
            sp = self._targetSysDict["amd64"].get_latest_action_dirpath()
            dstOsDir = os.path.join(osDir, "amd64")

            os.mkdir(dstOsDir)
            for fn in os.listdir(os.path.join(sp, "boot")):
                if fn.startswith("kernel-"):
                    shutil.copy(os.path.join(sp, "boot", fn), os.path.join(dstOsDir, "vmlinuz"))
                elif fn.startswith("initramfs-") and not fn.startswith("initramfs-files-"):
                    shutil.copy(os.path.join(sp, "boot", fn), os.path.join(dstOsDir, "initramfs.img"))
            shutil.copy(os.path.join(sp, "usr", "share", "memtest86+", "memtest.efi64"), dstOsDir)
            shutil.copy(os.path.join(sp, "usr", "share", "memtest86+", "memtest64.bios"), dstOsDir)
            Util.makeSquashedRootfsFiles(sp, dstOsDir, exclude=["boot/*"])

        # install target system files (arm64) into usb stick
        if "arm64" in self._targetSysDict:
            # FIXME
            assert False

        # install grub
        s = grub_install.Source(sp)
        t = grub_install.Target(grub_install.TargetType.MOUNTED_HDD_DEV, grub_install.TargetAccessMode.W, boot_mount_point=mp)
        if "amd64" in self._targetSysDict:
            t.install_platform(grub_install.PlatformType.X86_64_EFI, s, removable=True, update_nvram=False)
            t.install_platform(grub_install.PlatformType.I386_PC, s, rs_codes=False)
        if "arm64" in self._targetSysDict:
            # FIXME
            assert False
        t.install_data_files(s, locales="*", fonts="*", themes="*")

        # create grub.cfg
        with open(os.path.join(mp.mountpoint, "grub", "grub.cfg"), "w") as f:
            f.write("set default=0\n")
            f.write("set timeout=90\n")
            f.write("set gfxmode=auto\n")
            f.write("\n")

            f.write("insmod efi_gop\n")
            f.write("insmod efi_uga\n")
            f.write("insmod gfxterm\n")
            f.write("insmod all_video\n")
            f.write("insmod videotest\n")
            f.write("insmod videoinfo\n")
            f.write("terminal_output gfxterm\n")
            f.write("\n")

            if "amd64" in self._targetSysDict and "arm64" in self._targetSysDict:
                # FIXME
                assert False
            elif "amd64" in self._targetSysDict:
                f.write("menuentry \"Boot %s\" --class gnu-linux --class os {\n" % (self._diskName))
                f.write("    linux /os/amd64/vmlinuz\n")
                f.write("    initrd /os/amd64/initramfs.img\n")
                f.write("}\n")
            elif "arm64" in self._targetSysDict:
                # FIXME
                assert False
            else:
                assert False
            f.write("\n")

            f.write("menuentry \"Boot existing OS\" --class os {\n")
            f.write("    set root=(hd0)\n")
            f.write("    chainloader +1\n")
            f.write("}\n")
            f.write("\n")

            if "amd64" in self._targetSysDict and "arm64" in self._targetSysDict:
                # FIXME
                assert False
            elif "amd64" in self._targetSysDict:
                f.write('if [ "${grub_platform}" == "efi" ]; then\n')
                f.write("    menuentry \"Run Memtest86+ (RAM test)\" {\n")
                f.write('        chainloader "/os/amd64/memtest.efi64"\n')
                f.write("    }\n")
                f.write('elif [ "${grub_platform}" == "pc"]; then\n')
                f.write("    menuentry \"Run Memtest86+ (RAM test)\" {\n")
                f.write('        linux "/os/amd64/memtest64.bios"\n')
                f.write("    }\n")
                f.write('fi\n')
            elif "arm64" in self._targetSysDict:
                # FIXME
                assert False
            else:
                assert False
            f.write("\n")

            # menuentry "Hardware Information (HDT)" {
            #     linux /os/%ARCH%/hdt
            # }

            # Menu
            f.write("menuentry \"Restart\" {\n")
            f.write("    reboot\n")
            f.write("}\n")
            f.write("\n")

            # Menu
            f.write("menuentry \"Power Off\" {\n")
            f.write("    halt\n")
            f.write("}\n")

        # create livecd
        # FIXME: it sucks that genkernel's initrd requires this file
        with open(os.path.join(mp.mountpoint, "livecd"), "w") as f:
            f.write("")


class AssistDiskRemote:

    def __init__(self, hostname, port):
        self.hostname = hostname
        self.ctrlPort = port

        self.sock = None
        self.sslSock = None

    def connectAndInit(self):
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.connect((self.hostname, self.ctrlPort))

            ctx = SSL.Context(SSL.TLSv1_2_METHOD)
            ctx.use_certificate_file(FmConst.myCertFile)
            ctx.use_privatekey_file(FmConst.myPrivKeyFile)
            self.sslSock = SSL.Connection(ctx, self.sock)
            self.sslSock.set_connect_state()

            self._sendRequestObj({
                "command": "init",
                "hostname": socket.gethostname(),
                "cpu-arch": Util.getCpuArch(),
                "cpu-model": Util.getCpuModel(),
                "plugin": "gentoo",
            })
            resp = self._recvReponseObj()
            if "error" in resp:
                raise Exception(resp["error"])
        except:
            self.dispose()
            raise

    def dispose(self):
        assert self.asyncJobCount == 0

        if self.identityFile is None:
            os.unlink(self.identityFile)
            self.identityFile = None
        if self.cfgFile is None:
            os.unlink(self.cfgFile)
            self.cfgFile = None
        self.wSshPort = None
        self.wRsyncPort = None
        self.wCatFilePort = None

        if self.sslSock is not None:
            self.sslSock.close()
            self.sslSock = None
        if self.sock is not None:
            self.sock.close()
            self.sock = None

        self.ctrlPort = None
        self.hostname = None

    def createStorageLayout(self, layoutName):
        pass

    def sendSystemFiles(self, layoutName):
        pass
